"""
    SpectralMethods

Advanced spectral decomposition methods for background error covariance modeling in GSI.
This module implements sophisticated spectral analysis techniques including spherical 
harmonics, Fast Fourier Transforms, and eigenvalue decomposition methods used in 
operational atmospheric data assimilation.

The spectral approach provides an efficient representation of background error covariance
through eigenvalue decomposition:

```
B = Σᵢ λᵢ φᵢ φᵢᵀ = Φ Λ Φᵀ
```

Where λᵢ are eigenvalues, φᵢ are eigenvectors (often spherical harmonics), and the
decomposition enables efficient covariance operations and spectral truncation.

Key Capabilities:
- Spherical harmonic transforms (forward/inverse) for global grids
- Fast Fourier Transform integration for efficient computations
- Eigenvalue decomposition with spectral truncation
- Spectral filtering and covariance matrix factorization
- Support for both full and reduced-rank representations

Mathematical Foundation:
The spectral representation exploits the natural basis functions of atmospheric dynamics:
- Spherical harmonics Yₗᵐ(θ,φ) for global applications
- Fourier modes exp(ikx) for periodic boundaries
- Empirical Orthogonal Functions (EOFs) for data-driven bases

This enables efficient O(N log N) operations compared to O(N²) direct matrix methods.
"""
module SpectralMethods

using LinearAlgebra
using FFTW
using StaticArrays

export SpectralBasis, SphericalHarmonics, FourierBasis, EOFBasis
export SpectralDecomposition, TruncatedSpectralDecomposition
export forward_transform, inverse_transform, spectral_truncate
export compute_eigendecomposition, apply_spectral_filter
export spherical_harmonic_transform, associate_legendre_polynomials

"""
    AbstractSpectralBasis{T<:AbstractFloat}

Abstract base type for spectral basis function implementations.
"""
abstract type AbstractSpectralBasis{T<:AbstractFloat} end

"""
    SphericalHarmonics{T<:AbstractFloat}

Spherical harmonic basis functions for global atmospheric grids.

Mathematical basis:
```
Yₗᵐ(θ,φ) = √[(2l+1)(l-|m|)!/4π(l+|m|)!] * Pₗᵐ(cos θ) * exp(imφ)
```

# Fields
- `max_degree::Int`: Maximum spherical harmonic degree L
- `max_order::Int`: Maximum order M (≤ L)
- `truncation_type::Symbol`: :triangular or :rhomboidal truncation
- `normalization::Symbol`: :schmidt or :geodetic normalization
- `legendre_cache::Matrix{T}`: Cached associated Legendre polynomials
- `gaussian_weights::Vector{T}`: Gaussian quadrature weights
- `gaussian_latitudes::Vector{T}`: Gaussian quadrature latitudes
"""
struct SphericalHarmonics{T<:AbstractFloat} <: AbstractSpectralBasis{T}
    max_degree::Int
    max_order::Int
    truncation_type::Symbol
    normalization::Symbol
    legendre_cache::Matrix{T}
    gaussian_weights::Vector{T}
    gaussian_latitudes::Vector{T}
    
    function SphericalHarmonics{T}(max_degree::Int, max_order::Int = max_degree;
                                  truncation_type::Symbol = :triangular,
                                  normalization::Symbol = :schmidt,
                                  n_latitudes::Int = max_degree + 1) where T
        
        # Generate Gaussian quadrature points for integration
        lats, weights = gaussian_quadrature_latitudes(T, n_latitudes)
        
        # Pre-compute associated Legendre polynomials
        n_modes = triangular_truncation_size(max_degree, max_order)
        legendre_cache = Matrix{T}(undef, n_modes, length(lats))
        
        mode_idx = 1
        for l in 0:max_degree
            for m in 0:min(l, max_order)
                for (lat_idx, lat) in enumerate(lats)
                    legendre_cache[mode_idx, lat_idx] = 
                        associated_legendre_polynomial(l, m, cos(lat), normalization)
                end
                mode_idx += 1
            end
        end
        
        new{T}(max_degree, max_order, truncation_type, normalization,
               legendre_cache, weights, lats)
    end
end

SphericalHarmonics(max_degree::Int, args...; kwargs...) = 
    SphericalHarmonics{Float64}(max_degree, args...; kwargs...)

"""
    FourierBasis{T<:AbstractFloat}

Fourier basis functions for periodic boundary conditions.

# Fields
- `n_modes::Int`: Number of Fourier modes
- `domain_length::T`: Physical domain length
- `wavenumbers::Vector{T}`: Wavenumber array
- `fft_plan::FFTW.Plan`: Pre-computed FFT plan for efficiency
"""
struct FourierBasis{T<:AbstractFloat} <: AbstractSpectralBasis{T}
    n_modes::Int
    domain_length::T
    wavenumbers::Vector{T}
    fft_plan::FFTW.Plan
    
    function FourierBasis{T}(n_modes::Int, domain_length::T = T(2π)) where T
        # Construct wavenumber array
        wavenumbers = fftfreq(n_modes, 2π/domain_length * n_modes)
        
        # Create optimized FFT plan
        dummy_data = zeros(Complex{T}, n_modes)
        fft_plan = plan_fft(dummy_data)
        
        new{T}(n_modes, domain_length, wavenumbers, fft_plan)
    end
end

FourierBasis(n_modes::Int, domain_length::AbstractFloat = 2π) = 
    FourierBasis{typeof(domain_length)}(n_modes, domain_length)

"""
    EOFBasis{T<:AbstractFloat}

Empirical Orthogonal Function basis from data covariance analysis.

# Fields
- `eigenvectors::Matrix{T}`: EOF spatial patterns (columns are EOFs)
- `eigenvalues::Vector{T}`: Eigenvalues (explained variance)
- `cumulative_variance::Vector{T}`: Cumulative variance explained
- `mean_state::Vector{T}`: Mean state for anomaly computation
- `n_retained::Int`: Number of retained EOFs
"""
struct EOFBasis{T<:AbstractFloat} <: AbstractSpectralBasis{T}
    eigenvectors::Matrix{T}
    eigenvalues::Vector{T}
    cumulative_variance::Vector{T}
    mean_state::Vector{T}
    n_retained::Int
    
    function EOFBasis{T}(data_matrix::Matrix{T}, variance_threshold::T = T(0.95)) where T
        n_space, n_time = size(data_matrix)
        
        # Compute mean state and anomalies
        mean_state = vec(mean(data_matrix, dims=2))
        anomalies = data_matrix .- mean_state
        
        # Compute covariance matrix (spatial)
        covariance_matrix = (anomalies * anomalies') / (n_time - 1)
        
        # Eigenvalue decomposition
        eigenvals, eigenvecs = eigen(Symmetric(covariance_matrix))
        
        # Sort in descending order of eigenvalues
        sort_indices = sortperm(eigenvals, rev=true)
        eigenvals = eigenvals[sort_indices]
        eigenvecs = eigenvecs[:, sort_indices]
        
        # Compute cumulative variance
        total_variance = sum(eigenvals)
        cumulative_variance = cumsum(eigenvals) / total_variance
        
        # Determine number of modes to retain
        n_retained = findfirst(x -> x >= variance_threshold, cumulative_variance)
        n_retained = n_retained === nothing ? length(eigenvals) : n_retained
        
        new{T}(eigenvecs, eigenvals, cumulative_variance, mean_state, n_retained)
    end
end

"""
    SpectralDecomposition{T<:AbstractFloat, B<:AbstractSpectralBasis{T}}

Complete spectral decomposition of a covariance matrix using specified basis.

# Fields
- `basis::B`: Spectral basis functions
- `coefficients::Matrix{T}`: Spectral coefficients matrix
- `eigenvalues::Vector{T}`: Eigenvalues of decomposition
- `truncation_level::Int`: Number of retained modes
- `explained_variance::T`: Fraction of variance explained by retained modes
"""
struct SpectralDecomposition{T<:AbstractFloat, B<:AbstractSpectralBasis{T}}
    basis::B
    coefficients::Matrix{T}
    eigenvalues::Vector{T}
    truncation_level::Int
    explained_variance::T
end

"""
    forward_transform(basis::SphericalHarmonics{T}, field::Matrix{T}) where T

Forward spherical harmonic transform: physical → spectral space.

# Arguments
- `basis::SphericalHarmonics{T}`: Spherical harmonic basis
- `field::Matrix{T}`: Physical space field (longitude × latitude)

# Returns
- `Vector{Complex{T}}`: Spherical harmonic coefficients

# Mathematical Operation
```
aₗᵐ = ∫∫ f(θ,φ) Yₗᵐ*(θ,φ) sin(θ) dθ dφ
```
"""
function forward_transform(basis::SphericalHarmonics{T}, field::Matrix{T}) where T
    n_lon, n_lat = size(field)
    
    # Check grid compatibility
    if n_lat != length(basis.gaussian_latitudes)
        throw(DimensionMismatch("Field latitudes must match Gaussian grid"))
    end
    
    # Compute longitude Fourier transforms first
    lon_coeffs = fft(field, 1) / n_lon
    
    # Then compute latitude integrations using Gaussian quadrature
    n_modes = size(basis.legendre_cache, 1)
    spectral_coeffs = zeros(Complex{T}, n_modes)
    
    mode_idx = 1
    for l in 0:basis.max_degree
        for m in 0:min(l, basis.max_order)
            # Integrate over latitudes using Gaussian quadrature
            for (lat_idx, weight) in enumerate(basis.gaussian_weights)
                # Use appropriate longitude mode
                if m <= n_lon ÷ 2
                    lon_coeff = m == 0 ? real(lon_coeffs[1, lat_idx]) : 
                                       2 * real(lon_coeffs[m+1, lat_idx])
                else
                    lon_coeff = zero(T)
                end
                
                spectral_coeffs[mode_idx] += weight * 
                    basis.legendre_cache[mode_idx, lat_idx] * lon_coeff
            end
            mode_idx += 1
        end
    end
    
    return spectral_coeffs
end

"""
    inverse_transform(basis::SphericalHarmonics{T}, coefficients::Vector{Complex{T}}, 
                     grid_size::Tuple{Int,Int}) where T

Inverse spherical harmonic transform: spectral → physical space.

# Arguments
- `basis::SphericalHarmonics{T}`: Spherical harmonic basis
- `coefficients::Vector{Complex{T}}`: Spectral coefficients
- `grid_size::Tuple{Int,Int}`: Output grid size (n_lon, n_lat)

# Returns
- `Matrix{T}`: Physical space field

# Mathematical Operation
```
f(θ,φ) = Σₗ Σₘ aₗᵐ Yₗᵐ(θ,φ)
```
"""
function inverse_transform(basis::SphericalHarmonics{T}, 
                          coefficients::Vector{Complex{T}},
                          grid_size::Tuple{Int,Int}) where T
    n_lon, n_lat = grid_size
    
    if n_lat != length(basis.gaussian_latitudes)
        throw(DimensionMismatch("Grid latitudes must match Gaussian grid"))
    end
    
    field = zeros(T, n_lon, n_lat)
    
    # Reconstruct field from spectral coefficients
    mode_idx = 1
    for l in 0:basis.max_degree
        for m in 0:min(l, basis.max_order)
            coeff = coefficients[mode_idx]
            
            for lat_idx in 1:n_lat
                legendre_val = basis.legendre_cache[mode_idx, lat_idx]
                
                # Add contribution to all longitude points
                if m == 0
                    # Zonal mean (m=0) term
                    field[:, lat_idx] .+= real(coeff) * legendre_val
                else
                    # Non-zonal terms with longitude dependence
                    for lon_idx in 1:n_lon
                        lon = 2π * (lon_idx - 1) / n_lon
                        exp_factor = exp(im * m * lon)
                        field[lon_idx, lat_idx] += real(coeff * exp_factor * legendre_val)
                    end
                end
            end
            mode_idx += 1
        end
    end
    
    return field
end

"""
    forward_transform(basis::FourierBasis{T}, signal::Vector{T}) where T

Forward Fourier transform using pre-computed FFT plan.
"""
function forward_transform(basis::FourierBasis{T}, signal::Vector{T}) where T
    if length(signal) != basis.n_modes
        throw(DimensionMismatch("Signal length must match basis size"))
    end
    
    return basis.fft_plan * complex.(signal)
end

"""
    inverse_transform(basis::FourierBasis{T}, coefficients::Vector{Complex{T}}) where T

Inverse Fourier transform using pre-computed plan.
"""
function inverse_transform(basis::FourierBasis{T}, coefficients::Vector{Complex{T}}) where T
    return real(inv(basis.fft_plan) * coefficients)
end

"""
    compute_eigendecomposition(covariance_matrix::Matrix{T}, 
                              basis::AbstractSpectralBasis{T};
                              truncation_threshold::T = T(0.95)) where T

Compute eigenvalue decomposition in spectral space with optional truncation.
"""
function compute_eigendecomposition(covariance_matrix::Matrix{T}, 
                                   basis::AbstractSpectralBasis{T};
                                   truncation_threshold::T = T(0.95)) where T
    
    # Compute eigendecomposition
    eigenvals, eigenvecs = eigen(Symmetric(covariance_matrix))
    
    # Sort in descending order
    sort_indices = sortperm(eigenvals, rev=true)
    eigenvals = eigenvals[sort_indices]
    eigenvecs = eigenvecs[:, sort_indices]
    
    # Determine truncation level
    cumulative_variance = cumsum(eigenvals) / sum(eigenvals)
    truncation_level = findfirst(x -> x >= truncation_threshold, cumulative_variance)
    truncation_level = truncation_level === nothing ? length(eigenvals) : truncation_level
    
    explained_variance = cumulative_variance[truncation_level]
    
    return SpectralDecomposition(basis, eigenvecs, eigenvals, 
                               truncation_level, explained_variance)
end

"""
    spectral_truncate(decomposition::SpectralDecomposition{T}, n_modes::Int) where T

Truncate spectral decomposition to specified number of modes.
"""
function spectral_truncate(decomposition::SpectralDecomposition{T}, n_modes::Int) where T
    n_modes = min(n_modes, length(decomposition.eigenvalues))
    
    truncated_eigenvals = decomposition.eigenvalues[1:n_modes]
    truncated_eigenvecs = decomposition.coefficients[:, 1:n_modes]
    
    explained_variance = sum(truncated_eigenvals) / sum(decomposition.eigenvalues)
    
    return SpectralDecomposition(decomposition.basis, truncated_eigenvecs,
                               truncated_eigenvals, n_modes, explained_variance)
end

"""
    apply_spectral_filter(decomposition::SpectralDecomposition{T}, 
                         field::AbstractVector{T}) where T

Apply spectral filtering using truncated decomposition.
"""
function apply_spectral_filter(decomposition::SpectralDecomposition{T}, 
                              field::AbstractVector{T}) where T
    n_retain = decomposition.truncation_level
    
    # Project onto retained modes
    coefficients = decomposition.coefficients[:, 1:n_retain]' * field
    
    # Reconstruct with filtering
    filtered_field = coefficients' * coefficients
    
    return filtered_field
end

"""
    associated_legendre_polynomial(l::Int, m::Int, x::T, 
                                  normalization::Symbol = :schmidt) where T

Compute associated Legendre polynomial Pₗᵐ(x) with specified normalization.
"""
function associated_legendre_polynomial(l::Int, m::Int, x::T, 
                                       normalization::Symbol = :schmidt) where T
    if abs(m) > l
        return zero(T)
    end
    
    # Handle negative orders
    if m < 0
        m = -m
        sign_factor = (-1)^m
    else
        sign_factor = one(T)
    end
    
    # Compute using recurrence relations
    if m == 0
        # Legendre polynomial P_l(x)
        if l == 0
            result = one(T)
        elseif l == 1
            result = x
        else
            # Recurrence: (l+1)P_{l+1} = (2l+1)xP_l - lP_{l-1}
            p0, p1 = one(T), x
            for n in 2:l
                p_next = ((2n-1) * x * p1 - (n-1) * p0) / n
                p0, p1 = p1, p_next
            end
            result = p1
        end
    else
        # Associated Legendre polynomial
        # Start with P_m^m(x) = (-1)^m (2m-1)!! (1-x²)^{m/2}
        factor = (-1)^m
        for i in 1:m
            factor *= (2i - 1)
        end
        
        sqrt_term = (1 - x^2)^(m/2)
        
        if l == m
            result = factor * sqrt_term
        else
            # Use recurrence for higher l
            pmm = factor * sqrt_term
            pmm1 = x * (2m + 1) * pmm
            
            if l == m + 1
                result = pmm1
            else
                for n in (m+2):l
                    pmm2 = (x * (2n - 1) * pmm1 - (n + m - 1) * pmm) / (n - m)
                    pmm, pmm1 = pmm1, pmm2
                end
                result = pmm1
            end
        end
    end
    
    # Apply normalization
    if normalization == :schmidt
        # Schmidt semi-normalization
        if m == 0
            norm_factor = sqrt((2l + 1) / 2)
        else
            norm_factor = sqrt((2l + 1) * factorial(l - m) / factorial(l + m))
        end
        result *= norm_factor
    elseif normalization == :geodetic
        # Fully normalized (geodetic)
        norm_factor = sqrt((2l + 1) * factorial(l - m) / (2 * factorial(l + m)))
        result *= norm_factor
    end
    
    return sign_factor * result
end

"""
    gaussian_quadrature_latitudes(::Type{T}, n_points::Int) where T

Compute Gaussian quadrature points and weights for latitude integration.
"""
function gaussian_quadrature_latitudes(::Type{T}, n_points::Int) where T
    # Compute Gauss-Legendre quadrature points and weights
    # This is a simplified implementation - production would use optimized library
    
    latitudes = zeros(T, n_points)
    weights = zeros(T, n_points)
    
    # Symmetric points about zero
    for i in 1:n_points
        # Initial guess for roots
        theta = π * (i - 0.25) / (n_points + 0.5)
        x = cos(theta)
        
        # Newton-Raphson iteration to find roots of Legendre polynomial
        for iter in 1:10
            p1 = one(T)
            p2 = zero(T)
            
            # Evaluate Legendre polynomial and derivative
            for j in 1:n_points
                p3 = p2
                p2 = p1
                p1 = ((2j - 1) * x * p2 - (j - 1) * p3) / j
            end
            
            pp = n_points * (x * p1 - p2) / (x^2 - 1)
            x_new = x - p1 / pp
            
            if abs(x_new - x) < 1e-14
                break
            end
            x = x_new
        end
        
        latitudes[i] = acos(x)
        weights[i] = 2 / ((1 - x^2) * pp^2)
    end
    
    return latitudes, weights
end

"""
    triangular_truncation_size(max_degree::Int, max_order::Int)

Compute number of modes in triangular truncation.
"""
function triangular_truncation_size(max_degree::Int, max_order::Int)
    total = 0
    for l in 0:max_degree
        total += min(l, max_order) + 1
    end
    return total
end

end # module SpectralMethods